feat: optional AWS Lambda MicroVM stateful sandbox backend#16
feat: optional AWS Lambda MicroVM stateful sandbox backend#16danny-avila wants to merge 50 commits into
Conversation
Moves the sandbox execute POST from workers.ts into a pluggable HttpSandboxBackend behind getSandboxBackend(). Adds CODEAPI_SANDBOX_BACKEND config (http default) and a startup policy check that rejects lambda-microvm until that backend lands. No wire behavior change: the signed request body and headers pass through byte-identical, and axios errors are rethrown untouched so the worker's abort/timeout/sandbox-error mapping is unchanged.
… launch) Adds runtime_session_hint to /exec (validated, optional) and derives rt_<sha256(namespace,user,hint)> server-side so a client hint can never collide across tenants. The id rides JobData into the sandbox backend context; stateless mode (default) derives nothing and enqueues byte-identical job data. The registry maps runtime_session_id -> MicroVM record in Redis with the replay-state lock discipline: SET NX PX mutex, CAS-delete release, token-fenced record writes/removals, monotonic generation counter for launch fencing, and a last-seen zset for the idle sweeper. No consumers yet - the Lambda backend lands behind CODEAPI_SANDBOX_BACKEND.
Adds the LambdaMicrovmClient interface plus AwsLambdaMicrovmClient (@aws-sdk/client-lambda-microvms, isolated to lambda-client-aws.ts and absent from http-only bundles), a transport-free in-memory fake for bun tests, and Redis-backed per-second token buckets with poison backoff for the account-wide control-plane TPS limits. Mapping notes from the SDK typings: RunMicrovm takes imageIdentifier + imageVersion, connector ARN arrays, native idlePolicy (auto-suspend / auto-terminate / auto-resume), runHookPayload, and a clientToken idempotency key; auth tokens come back as an X-aws-proxy-auth header map with minute-granularity expiry (max 60).
…kend
Adds the AWS Lambda MicroVM execution path (opt-in, config-gated behind
CODEAPI_SANDBOX_BACKEND=lambda-microvm; default http is unchanged):
- api/Dockerfile lambda-microvm-runner target: the existing sandbox-runner
packaged for Lambda MicroVMs (arm64, port 8080, /pkgs baked, no libkrun
launcher since the MicroVM is the VM boundary).
- api lifecycle hooks at /aws/lambda-microvms/runtime/v1/{ready,run,resume,
suspend,terminate}: no-op acks in Phase 1-2, /run captures the per-VM
runHookPayload (Phase 3 checkpoint attachment point).
- LambdaMicrovmSandboxBackend (stateless mode): run -> poll RUNNING ->
mint X-aws-proxy-auth token -> health -> proxy /api/v2/execute ->
terminate (incl. terminate-on-abort); throttle-aware, metrics-wired.
- Startup policy rules (blocking-PTC reject, image-ARN required, hardened
egress-connector required, token-TTL cap, no shell in prod).
- entrypoint raises RLIMIT_NOFILE hard cap to 65536: the AL2023 MicroVM
base caps it at 1024 (below the 2048 sandbox default), which made every
in-guest nsjail job fail setrlimit with EPERM. Verified on live AWS.
Live-verified on AWS Lambda MicroVMs (us-east-1): image builds and
snapshots, nsjail runs with additionalOsCapabilities ALL, bash+python
execute (code 0), and suspend/resume preserves process+memory state
with ~1.2s auto-resume-on-traffic.
Turns the semi-stateless runner stateful when a VM is bound to a runtime
session: instead of a fresh /mnt/data workspace per /execute, the VM
reuses ONE persistent workspace and one pinned UID across calls, so
files, installed packages, and chDB dirs survive between tool calls.
Modular and off by default: gated by TWO independent locks, so the
legacy fresh-per-job path is byte-identical when either is absent.
1. Image-level SANDBOX_SESSION_WORKSPACE_ENABLED — set only in the
lambda-microvm-runner target; the K8s image cannot enter session
mode regardless of any payload.
2. Per-launch /run runHookPayload {runtime_session_id, session_workspace}
— the control plane opts a specific VM in.
Mechanism:
- session-workspace.ts: process singleton bound by the /run hook,
unbound by /terminate; holds the pinned UID + output-diff and
priming-dedup state.
- workspace-isolation.ts: ensureSessionWorkspace (stable id, contents
preserved, reaper-protected) + resetSessionWorkspace teardown.
- Job branches only at three seams — prime (reuse workspace+UID, skip
re-downloading unchanged inputs), walkDir file surfacing (skip prior
outputs unchanged by size+mtime — output diffing), cleanup (keep
workspace+UID for the session).
Verified end-to-end in the arm64 runner container: fresh mode wipes
between calls; session mode persists files across calls (incl. across
languages) and accumulates; /terminate wipes and rebind gives a clean
slate. 274 api tests pass (fresh path unchanged).
Connects the proven runner session workspace to the product: when a
runtime session id is present and the mode is affinity/strict, the
Lambda backend finds-or-launches ONE warm VM per runtime_session_id via
the registry, delivers the /run runHookPayload
{runtime_session_id, session_workspace:true} that activates the runner's
persistent workspace, and reuses that VM across executes. AWS idlePolicy
(autoResume) suspends the VM when idle and wakes it on the next request,
so there is no explicit resume in the hot path.
- Serializes per session on the registry lock; strict contention -> 409
(publicExecutionFailure maps RUNTIME_SESSION_BUSY), affinity contention
-> correct stateless one-shot fallback (files always ride the payload).
- Generation-fenced launch: a fenced worker terminates its orphan VM.
- Stateless path unchanged (one VM per exec, terminate after).
- Lifts the stateless-only startup policy; adds session/fallback/lock
metrics. 345 service tests pass, incl. reuse/serialize/fallback/fence.
… across expiry
Makes an expiring/evicted MicroVM's state survive a relaunch — the
difference between real statefulness and just re-implementing the
existing file-ref system. The file-ref path only brings back files
surfaced as CodeEnvRefs; checkpoint/restore brings back the WHOLE
workspace: pip-installed packages, venvs, chDB dirs, caches, and files
with unsupported extensions.
Two runner endpoints (session-mode only, 409 otherwise so the legacy
runner exposes nothing new):
- GET /api/v2/session/checkpoint streams tar.gz of the workspace
- POST /api/v2/session/restore replaces the workspace from one,
re-owned to the session's pinned UID
Control-plane driven: the orchestrator pulls the checkpoint over the
authed proxy and owns the S3 write, so the untrusted VM never gets S3
credentials (matches the report's checkpoint-capability security model).
Verified end-to-end with two containers simulating VM expiry: VM1 builds
a python module tree + a 2KB unsupported-extension binary, is
terminated; a fresh VM2 shows the state absent (all file-refs give you),
then after restore imports the module (greet()=42) and reads the binary
(2048 bytes) — full workspace continuity across a VM swap. 276 api
tests pass.
Closes the 8h-rollover / eviction story so perceived statefulness is
automatic instead of control-plane-by-hand. After each successful
session exec (lock still held), pull the workspace tar from the warm VM
and store it to S3 under a deterministic key (rtsx-checkpoints/<id>.tar.gz)
so recovery survives even registry loss; record the pointer under the
same fenced write. On relaunch, a fresh session VM restores its
predecessor's checkpoint before the first exec, making an expired/evicted
VM invisible.
- Coverage is complete and tear-free: the workspace only mutates during
an exec and execs serialize on the session lock, so the post-exec
checkpoint always captures the latest state; a busy lock means a newer
exec will checkpoint instead.
- Never fatal: a missed checkpoint degrades to file-ref recovery, a
failed restore to a fresh workspace ('relaunched must be correct').
- Off => warm reuse still works, cross-VM restore falls back to file
refs. CheckpointStore injected (Minio prod / Memory in tests); byte
cap + timeout bound the transfer; checkpoint/restore/bytes metrics.
354 service tests pass incl. checkpoint-after-exec, restore-before-first
-exec ordering, no-restore-on-reuse, disabled skip, failure non-fatal.
…eader Deliver session mode per /execute request instead of the /run lifecycle hook. Lambda image build hooks only route on the snapshot-compatible base container image, and enabling any runtime hook forces the /ready build hook (which never reaches a stock container listener), so hookless image builds are the reliable path. The runner binds its persistent workspace from the header; the backend stamps it on the proxied execute in session mode and drops runHookPayload from RunMicrovm. Verified on a real hookless MicroVM.
|
|
…lper Operator guide for the optional AWS Lambda MicroVM backend: the cross-repo picture, from-scratch AWS setup, a full config reference, operating modes, alternative AWS methods (base image, checkpoint store, egress, quota), the PTC replay/blocking distinction, and the hard-won runbook gotchas. - docs/lambda-microvm/terraform: prerequisites module (checkpoint + artifact buckets, build + logging-only execution roles with the sts:TagSession / logs:* trust the build needs, CloudWatch log groups, checkpoint access policy). terraform validate + fmt clean. - service/scripts/create-microvm-image.ts: guaranteed-correct hookless CreateMicrovmImage helper (ALL caps + cgroupv2 off baked in), the one provisioning step Terraform can't own.
…, +)
Fixes from the Macroscope review pass:
- registry: derive RUNTIME_SESSION_LOCK_TTL_MS from the actual launch/health/
execute/checkpoint budgets (was a 60s placeholder, could expire mid-work at
defaults); guard readRuntimeSessionRecord JSON.parse so a corrupt key reads
as missing instead of wedging the session.
- checkpoint: fence (fenced record write) BEFORE the deterministic-key S3 put
so a lock-expired caller can't clobber a newer blob; cap restore size against
maxBytes before buffering.
- session-checkpoint: lchown + never follow symlinks when chowning a restored
(untrusted) checkpoint, so a symlinked entry can't re-own files outside the
workspace.
- lambda-client: throw when a MicroVM response omits microvmId instead of
returning '' (a partial RunMicrovm response would otherwise orphan a billable
VM behind getMicrovm('')/terminateMicrovm('')).
- router: skip runtime_session_hint validation in stateless mode (the field is
ignored there, so a malformed hint must not 400).
- v2/session: only run in the persistent workspace when THIS request carried a
valid X-Runtime-Session-Id, so a headerless request never inherits a prior
session's workspace/UID.
- entrypoint: only ever raise RLIMIT_NOFILE (guard so it never clamps a higher
host default down to 65536).
- secure-startup: warn (don't silently no-op) on affinity+http.
Deferred with reasons in the PR: zset-orphan (no consumer until the sweeper
PR), bindSessionWorkspace rebind race (prevented by the 1-VM-1-session
invariant), releaseLock swallow (TTL self-heals), /run applyRunHook (inert in
the hookless design), startupApiOnly policy placement (split-deploy design Q).
|
Addressed the review findings in Fixed:
Deferred, with reasoning:
|
|
@codex review |
- terraform: artifact + checkpoint buckets use SSE-S3 (AES256) so the s3:GetObject build role and the checkpoint access policy need no kms:Decrypt grant (SSE-KMS would AccessDenied on read). - terraform: build + runtime log policies grant both the log-group ARN and its stream form — stream-level actions (CreateLogStream/PutLogEvents) don't match the group ARN, which would fail builds with an empty stateReason. - terraform: validate artifact_bucket_name is non-empty when reusing an existing bucket (else the policy resolves to arn:aws:s3:::/*); bump required_version to >= 1.9 for cross-variable validation. - create-microvm-image.ts: cap the poll loop with a deadline (default 30m, MICROVM_BUILD_DEADLINE_MINUTES) and exit non-zero, so a wedged CREATING build can't hang a provisioning job forever. - README: add MINIO_PORT=443 to the S3 example (client defaults to 9000); scope the teardown sweep to VMs from this image's ARN so it can't terminate unrelated MicroVMs in a shared account.
|
@codex review |
- terraform: include region in the artifact + checkpoint bucket names (S3 names are global, so a same-prefix second-region apply would collide). - docs: correct the checkpoint credential guidance. The MinIO-compatible client reads only static MINIO_ACCESS_KEY/SECRET and does not load task-role/IRSA creds, so attaching checkpoint_access_policy_arn to a task role alone does not work. Point operators at create_checkpoint_access_user (static keys) and note IRSA-aware loading as a follow-up. Fixed across variables.tf, outputs.tf, README.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8e2da55797
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| try { | ||
| const existing = await readRuntimeSessionRecord(runtimeSessionId); | ||
| const vm = await this.findOrLaunchSession(client, ctx, runtimeSessionId, existing, lockToken); | ||
| const result = await this.proxyExecute(client, vm, req, ctx, runtimeSessionId); |
There was a problem hiding this comment.
Fixed in daa43b0. New executeOnSessionVm helper: on a reuse failure where the runner is unreachable (health/connection error — e.g. idlePolicy auto-terminated the suspended VM), it terminates the VM and removeRuntimeSession, so the next call relaunches + restores instead of reusing a dead VM still recorded RUNNING. A live-but-non-200 response leaves the warm VM and its record intact.
| } | ||
| this.sessionFiles.push(fileData); | ||
| this.generatedFiles.push({ id: newId, name: relativePath, path: fullPath }); | ||
| if (this.session) this.session.markSurfaced(relativePath, outputSignature); |
There was a problem hiding this comment.
Fixed in daa43b0. The surfaced-mark is now recorded per-file (id → name+signature) during the scan and committed to session.markSurfaced only for files that upload successfully, in uploadGeneratedFiles. A dropped upload leaves the file eligible to surface next turn.
| cpuConfigurations: [{ architecture: 'ARM_64' as const }], | ||
| resources: [{ minimumMemoryInMiB: memory }], | ||
| additionalOsCapabilities: ['ALL' as const], | ||
| environmentVariables: { SANDBOX_USE_CGROUPV2: 'false' }, |
There was a problem hiding this comment.
Fixed in daa43b0. create-microvm-image now takes --env-json / MICROVM_IMAGE_ENV_JSON, merged over SANDBOX_USE_CGROUPV2, so operators bake FILE_SERVER_URL / EGRESS_GATEWAY_URL / manifest env into the image at build time. README updated with the required keys and the "invalid /sessions/... URLs" rationale.
| const now = Date.now(); | ||
| const settled = await readRuntimeSessionRecord(runtimeSessionId); | ||
| const nextRecord = settled | ||
| ? await this.checkpointUnderLock(client, settled, runtimeSessionId, now, lockToken) |
There was a problem hiding this comment.
Deferred — valid, but the fix is a design choice. checkpointUnderLock is already best-effort (never throws), so a checkpoint failure doesn't lose the result, but its latency still counts against JOB_TIMEOUT. Decoupling it (complete the job then checkpoint, or include the checkpoint budget in the public wait) needs a call on where checkpointing sits relative to job completion — raised with the maintainer. Not a correctness issue.
| for await (const chunk of stream) { | ||
| chunks.push(chunk as Buffer); | ||
| } | ||
| return Buffer.concat(chunks); |
There was a problem hiding this comment.
Fixed in daa43b0. CheckpointStore.get now takes maxBytes and stats the S3 object size before downloading (throws CheckpointTooLargeError if over), so we never Buffer.concat an oversized object. Removed the now-redundant post-buffer guard in restoreSession; added a MemoryCheckpointStore oversize test.
- Bind the session on checkpoint/restore (F1): the hookless path runs /session/restore before the first /execute, so the runner had nothing bound and every real restore 409'd, losing checkpoint state across VM expiry. The runner now binds from X-Runtime-Session-Id in the checkpoint/restore routes, and the backend sends that header on both proxied calls. - Clear/terminate a session VM on reuse failure or abort (F2/F5): on a dead reused VM (idlePolicy auto-terminated) or an aborted execute (runner keeps NsJail alive after the socket closes), terminate the VM and drop the record so the next call relaunches + restores instead of reusing a dead-or-dirty VM. A plain non-200 from a live runner leaves the warm VM intact. - Enforce checkpoint size before buffering (F7): CheckpointStore.get takes maxBytes and stats the S3 object first, so an oversized/stray checkpoint can't OOM the worker before the (now-removed, too-late) post-buffer guard. - Mark session outputs surfaced only after upload (F3): a dropped upload no longer permanently suppresses an unchanged file next turn. - Bake runner file/egress/manifest env into the image (F4): create-microvm-image takes --env-json so images can reach FILE_SERVER_URL / EGRESS_GATEWAY_URL instead of building invalid /sessions/... URLs. Deferred (P2, design choice): don't count checkpoint time in the client job timeout — decoupling checkpoint from the response path needs a call on where it sits relative to job completion; raised with the maintainer.
| try { | ||
| const existing = await readRuntimeSessionRecord(runtimeSessionId); | ||
| const vm = await this.findOrLaunchSession(client, ctx, runtimeSessionId, existing, lockToken); | ||
| const result = await this.proxyExecute(client, vm, req, ctx, runtimeSessionId); |
There was a problem hiding this comment.
Fixed in daa43b0. New executeOnSessionVm helper: on a reuse failure where the runner is unreachable (health/connection error — e.g. idlePolicy auto-terminated the suspended VM), it terminates the VM and removeRuntimeSession, so the next call relaunches + restores instead of reusing a dead VM still recorded RUNNING. A live-but-non-200 response leaves the warm VM and its record intact.
| } | ||
| this.sessionFiles.push(fileData); | ||
| this.generatedFiles.push({ id: newId, name: relativePath, path: fullPath }); | ||
| if (this.session) this.session.markSurfaced(relativePath, outputSignature); |
There was a problem hiding this comment.
Fixed in daa43b0. The surfaced-mark is now recorded per-file (id → name+signature) during the scan and committed to session.markSurfaced only for files that upload successfully, in uploadGeneratedFiles. A dropped upload leaves the file eligible to surface next turn.
| cpuConfigurations: [{ architecture: 'ARM_64' as const }], | ||
| resources: [{ minimumMemoryInMiB: memory }], | ||
| additionalOsCapabilities: ['ALL' as const], | ||
| environmentVariables: { SANDBOX_USE_CGROUPV2: 'false' }, |
There was a problem hiding this comment.
Fixed in daa43b0. create-microvm-image now takes --env-json / MICROVM_IMAGE_ENV_JSON, merged over SANDBOX_USE_CGROUPV2, so operators bake FILE_SERVER_URL / EGRESS_GATEWAY_URL / manifest env into the image at build time. README updated with the required keys and the "invalid /sessions/... URLs" rationale.
| const now = Date.now(); | ||
| const settled = await readRuntimeSessionRecord(runtimeSessionId); | ||
| const nextRecord = settled | ||
| ? await this.checkpointUnderLock(client, settled, runtimeSessionId, now, lockToken) |
There was a problem hiding this comment.
Deferred — valid, but the fix is a design choice. checkpointUnderLock is already best-effort (never throws), so a checkpoint failure doesn't lose the result, but its latency still counts against JOB_TIMEOUT. Decoupling it (complete the job then checkpoint, or include the checkpoint budget in the public wait) needs a call on where checkpointing sits relative to job completion — raised with the maintainer. Not a correctness issue.
| for await (const chunk of stream) { | ||
| chunks.push(chunk as Buffer); | ||
| } | ||
| return Buffer.concat(chunks); |
There was a problem hiding this comment.
Fixed in daa43b0. CheckpointStore.get now takes maxBytes and stats the S3 object size before downloading (throws CheckpointTooLargeError if over), so we never Buffer.concat an oversized object. Removed the now-redundant post-buffer guard in restoreSession; added a MemoryCheckpointStore oversize test.
The two ends computed the SAME contract differently — the runner joined (storage session, id) with a literal NUL byte, the control plane with a space — so every pushed object landed under a name no lookup could ever produce. Both unit suites passed because each was self-consistent; the route test and the priming test used one side only. A live execution reported 'No such file' for files that had been delivered successfully. Both now use an explicit \u0000 escape (no raw NUL in source, which also made grep treat the file as binary), and both suites assert the same golden vector so the contract cannot silently drift again.
The cache is keyed by OBJECT, but the pushed metadata carried the file server's filename and the runner emitted it as Content-Disposition. Every ref for that object therefore resolved to one name, so requesting the same object at a second path wrote it to the FIRST path instead — overwriting a file the sandbox had edited in an earlier turn. Live testing caught it; the unit suites had encoded the wrong behaviour as expected. Metadata now carries only object-level facts (read-only), and priming falls back to each ref's requested name. Covered by a prime test that serves one cached object to two destinations.
|
@codex review |
|
Macroscope has since reviewed this pull request. An earlier review was skipped by a cost limit; a review has now completed, so that notice no longer applies. |
|
@codex review |
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d8d9f110bf
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
@codex review |
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d00b13b863
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
@codex review |
|
@codex review |
|
@codex review |
|
@codex review |
|
Codex Review: Didn't find any major issues. Delightful! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
What
Adds AWS Lambda MicroVMs as an optional, config-gated stateful execution backend for the Code Interpreter, giving perceived-indefinite statefulness (a warm per-session workspace plus checkpoint/restore across the VM's 8h lifetime) without changing the legacy semi-stateless HTTP path.
Design
SandboxBackendseam is extracted from the worker dispatch. Thehttpbackend is byte-identical to today.CODEAPI_SANDBOX_BACKEND=lambda-microvmopts in, and@aws-sdk/*is lazily imported so http deployments never load it.runtime_session_id = hash(storageNamespace, canonicalUserId, hint)plus a Redis registry (SET NX locks, Lua CAS release, generation fencing) pins one VM per session.CODEAPI_RUNTIME_SESSION_MODE=stateless|affinity|strict./mnt/dataand pinned UID across calls (prime dedup and output diffing), gated bySANDBOX_SESSION_WORKSPACE_ENABLED(Lambda image only) plus a per-request signal.idlePolicy(autoResume) handles idle suspend/resume, so the sweeper shrinks to registry hygiene.Hookless session binding
Session mode is delivered per request via an
X-Runtime-Session-Idheader rather than a/runlifecycle hook. Lambda's image build hooks only route on the snapshot-compatible base container image, and enabling any runtime hook forces the/readybuild hook, which never reaches a stock container's listener (builds then fail at the ready timeout). Hookless image builds are reliable; checkpoint/restore already runs over plain HTTP endpoints, and idle suspend/resume is native.Proven live
Real MicroVM: launch ~3s, suspend/resume ~2.2s with process continuity, auto-resume ~1.2s. Two-turn E2E on a real hookless MicroVM: with the header,
/mnt/datapersists across separate/executecalls (42 read back); without it, fresh-per-job (ABSENT).Tests
~630 api and service tests (ioredis-mock, fake AWS client, and Bun.serve fakes).